// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/substrait/substrait_to_pollux_expr.h>
#include <pollux/substrait/type_utils.h>
#include <pollux/vector/flat_vector.h>

using namespace kumo::pollux;

namespace {
    // Get values for the different supported types.
    template<typename T>
    T getLiteralValue(const ::substrait::Expression::Literal & /* literal */) {
        POLLUX_NYI();
    }

    template<>
    int8_t getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return static_cast<int8_t>(literal.i8());
    }

    template<>
    int16_t getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return static_cast<int16_t>(literal.i16());
    }

    template<>
    int32_t getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return literal.i32();
    }

    template<>
    int64_t getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return literal.i64();
    }

    template<>
    double getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return literal.fp64();
    }

    template<>
    float getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return literal.fp32();
    }

    template<>
    bool getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return literal.boolean();
    }

    template<>
    uint32_t getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return literal.i32();
    }

    template<>
    Timestamp getLiteralValue(const ::substrait::Expression::Literal &literal) {
        return Timestamp::fromMicros(literal.timestamp());
    }

    ArrayVectorPtr make_array_vector(const VectorPtr &elements) {
        BufferPtr offsets = allocateOffsets(1, elements->pool());
        BufferPtr sizes = allocateOffsets(1, elements->pool());
        sizes->asMutable<vector_size_t>()[0] = elements->size();

        return std::make_shared<ArrayVector>(
            elements->pool(),
            ARRAY(elements->type()),
            nullptr,
            1,
            offsets,
            sizes,
            elements);
    }

    ArrayVectorPtr makeEmptyArrayVector(memory::MemoryPool *pool) {
        BufferPtr offsets = allocateOffsets(1, pool);
        BufferPtr sizes = allocateOffsets(1, pool);
        return std::make_shared<ArrayVector>(
            pool, ARRAY(UNKNOWN()), nullptr, 1, offsets, sizes, nullptr);
    }

    template<typename T>
    void setLiteralValue(
        const ::substrait::Expression::Literal &literal,
        FlatVector<T> *vector,
        vector_size_t index) {
        if (literal.has_null()) {
            vector->set_null(index, true);
        } else if constexpr (std::is_same_v<T, StringView>) {
            if (literal.has_string()) {
                vector->set(index, StringView(literal.string()));
            } else if (literal.has_var_char()) {
                vector->set(index, StringView(literal.var_char().value()));
            } else {
                POLLUX_FAIL("Unexpected string literal");
            }
        } else if (vector->type()->isDate()) {
            auto dateVector = vector->template as_flat_vector<int32_t>();
            dateVector->set(index, static_cast<int32_t>(literal.date()));
        } else {
            vector->set(index, getLiteralValue<T>(literal));
        }
    }

    template<TypeKind kind>
    VectorPtr constructFlatVector(
        const ::substrait::Expression::Literal &listLiteral,
        const vector_size_t size,
        const TypePtr &type,
        memory::MemoryPool *pool) {
        POLLUX_CHECK(type->isPrimitiveType());
        auto vector = BaseVector::create(type, size, pool);
        using T = typename TypeTraits<kind>::NativeType;
        auto flat_vector = vector->as<FlatVector<T> >();

        vector_size_t index = 0;
        for (auto child: listLiteral.list().values()) {
            setLiteralValue(child, flat_vector, index++);
        }
        return vector;
    }

    /// Whether null will be returned on cast failure.
    bool isNullOnFailure(
        ::substrait::Expression::Cast::FailureBehavior failureBehavior) {
        switch (failureBehavior) {
            case ::substrait::
            Expression_Cast_FailureBehavior_FAILURE_BEHAVIOR_UNSPECIFIED:
            case ::substrait::
            Expression_Cast_FailureBehavior_FAILURE_BEHAVIOR_THROW_EXCEPTION:
                return false;
            case ::substrait::
            Expression_Cast_FailureBehavior_FAILURE_BEHAVIOR_RETURN_NULL:
                return true;
            default:
                POLLUX_NYI(
                    "The given failure behavior is NOT supported: '{}'", failureBehavior);
        }
    }
} // namespace

namespace kumo::pollux::substrait {
    std::shared_ptr<const core::FieldAccessTypedExpr>
    SubstraitPolluxExprConverter::toPolluxExpr(
        const ::substrait::Expression::FieldReference &substraitField,
        const RowTypePtr &inputType) {
        auto typeCase = substraitField.reference_type_case();
        switch (typeCase) {
            case ::substrait::Expression::FieldReference::ReferenceTypeCase::
            kDirectReference: {
                const auto &directRef = substraitField.direct_reference();
                int32_t colIdx = substraitParser_.parseReferenceSegment(directRef);
                const auto &inputNames = inputType->names();
                const int64_t inputSize = inputNames.size();
                if (colIdx <= inputSize) {
                    const auto &inputTypes = inputType->children();
                    // Convert type to row.
                    return std::make_shared<core::FieldAccessTypedExpr>(
                        inputTypes[colIdx],
                        std::make_shared<core::InputTypedExpr>(inputTypes[colIdx]),
                        inputNames[colIdx]);
                } else {
                    POLLUX_FAIL("Missing the column with id '{}' .", colIdx);
                }
            }
            default:
                POLLUX_NYI(
                    "Substrait conversion not supported for Reference '{}'", typeCase);
        }
    }

    core::TypedExprPtr SubstraitPolluxExprConverter::toPolluxExpr(
        const ::substrait::Expression::ScalarFunction &substraitFunc,
        const RowTypePtr &inputType) {
        std::vector<core::TypedExprPtr> params;
        params.reserve(substraitFunc.arguments().size());
        for (const auto &sArg: substraitFunc.arguments()) {
            params.emplace_back(toPolluxExpr(sArg.value(), inputType));
        }
        const auto &polluxFunction = substraitParser_.findPolluxFunction(
            functionMap_, substraitFunc.function_reference());
        return std::make_shared<const core::CallTypedExpr>(
            substraitParser_.parseType(substraitFunc.output_type()),
            std::move(params),
            polluxFunction);
    }

    std::shared_ptr<const core::ConstantTypedExpr>
    SubstraitPolluxExprConverter::toPolluxExpr(
        const ::substrait::Expression::Literal &substraitLit) {
        auto typeCase = substraitLit.literal_type_case();
        switch (typeCase) {
            case ::substrait::Expression_Literal::LiteralTypeCase::kBoolean:
                return std::make_shared<core::ConstantTypedExpr>(
                    BOOLEAN(), variant(substraitLit.boolean()));
            case ::substrait::Expression_Literal::LiteralTypeCase::kI8:
                // SubstraitLit.i8() will return int32, so we need this type conversion.
                return std::make_shared<core::ConstantTypedExpr>(
                    TINYINT(), variant(static_cast<int8_t>(substraitLit.i8())));
            case ::substrait::Expression_Literal::LiteralTypeCase::kI16:
                // SubstraitLit.i16() will return int32, so we need this type conversion.
                return std::make_shared<core::ConstantTypedExpr>(
                    SMALLINT(), variant(static_cast<int16_t>(substraitLit.i16())));
            case ::substrait::Expression_Literal::LiteralTypeCase::kI32:
                return std::make_shared<core::ConstantTypedExpr>(
                    INTEGER(), variant(substraitLit.i32()));
            case ::substrait::Expression_Literal::LiteralTypeCase::kFp32:
                return std::make_shared<core::ConstantTypedExpr>(
                    REAL(), variant(substraitLit.fp32()));
            case ::substrait::Expression_Literal::LiteralTypeCase::kI64:
                return std::make_shared<core::ConstantTypedExpr>(
                    BIGINT(), variant(substraitLit.i64()));
            case ::substrait::Expression_Literal::LiteralTypeCase::kFp64:
                return std::make_shared<core::ConstantTypedExpr>(
                    DOUBLE(), variant(substraitLit.fp64()));
            case ::substrait::Expression_Literal::LiteralTypeCase::kString:
                return std::make_shared<core::ConstantTypedExpr>(
                    VARCHAR(), variant(substraitLit.string()));
            case ::substrait::Expression_Literal::LiteralTypeCase::kNull: {
                auto polluxType = substraitParser_.parseType(substraitLit.null());
                return std::make_shared<core::ConstantTypedExpr>(
                    polluxType, variant::null(polluxType->kind()));
            }
            case ::substrait::Expression_Literal::LiteralTypeCase::kVarChar:
                return std::make_shared<core::ConstantTypedExpr>(
                    VARCHAR(), variant(substraitLit.var_char().value()));
            case ::substrait::Expression_Literal::LiteralTypeCase::kList: {
                auto constant_vector =
                        BaseVector::wrap_in_constant(1, 0, literalsToArrayVector(substraitLit));
                return std::make_shared<const core::ConstantTypedExpr>(constant_vector);
            }
            case ::substrait::Expression_Literal::LiteralTypeCase::kDate:
                return std::make_shared<core::ConstantTypedExpr>(
                    DATE(), variant(static_cast<int32_t>(substraitLit.date())));
            default:
                POLLUX_NYI(
                    "Substrait conversion not supported for type case '{}'", typeCase);
        }
    }

    ArrayVectorPtr SubstraitPolluxExprConverter::literalsToArrayVector(
        const ::substrait::Expression::Literal &listLiteral) {
        auto childSize = listLiteral.list().values().size();
        if (childSize == 0) {
            return makeEmptyArrayVector(pool_);
        }
        auto typeCase = listLiteral.list().values(0).literal_type_case();
        switch (typeCase) {
            case ::substrait::Expression_Literal::LiteralTypeCase::kBoolean:
                return make_array_vector(constructFlatVector<TypeKind::BOOLEAN>(
                    listLiteral, childSize, BOOLEAN(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kI8:
                return make_array_vector(constructFlatVector<TypeKind::TINYINT>(
                    listLiteral, childSize, TINYINT(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kI16:
                return make_array_vector(constructFlatVector<TypeKind::SMALLINT>(
                    listLiteral, childSize, SMALLINT(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kI32:
                return make_array_vector(constructFlatVector<TypeKind::INTEGER>(
                    listLiteral, childSize, INTEGER(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kFp32:
                return make_array_vector(constructFlatVector<TypeKind::REAL>(
                    listLiteral, childSize, REAL(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kI64:
                return make_array_vector(constructFlatVector<TypeKind::BIGINT>(
                    listLiteral, childSize, BIGINT(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kFp64:
                return make_array_vector(constructFlatVector<TypeKind::DOUBLE>(
                    listLiteral, childSize, DOUBLE(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kString:
            case ::substrait::Expression_Literal::LiteralTypeCase::kVarChar:
                return make_array_vector(constructFlatVector<TypeKind::VARCHAR>(
                    listLiteral, childSize, VARCHAR(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kNull: {
                auto polluxType = substraitParser_.parseType(listLiteral.null());
                auto kind = polluxType->kind();
                return make_array_vector(POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH(
                        constructFlatVector, kind, listLiteral, childSize, polluxType, pool_));
            }
            case ::substrait::Expression_Literal::LiteralTypeCase::kDate:
                return make_array_vector(constructFlatVector<TypeKind::INTEGER>(
                    listLiteral, childSize, DATE(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kTimestamp:
                return make_array_vector(constructFlatVector<TypeKind::TIMESTAMP>(
                    listLiteral, childSize, TIMESTAMP(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kIntervalDayToSecond:
                return make_array_vector(constructFlatVector<TypeKind::BIGINT>(
                    listLiteral, childSize, INTERVAL_DAY_TIME(), pool_));
            case ::substrait::Expression_Literal::LiteralTypeCase::kList: {
                VectorPtr elements;
                for (auto it: listLiteral.list().values()) {
                    auto v = literalsToArrayVector(it);
                    if (!elements) {
                        elements = v;
                    } else {
                        elements->append(v.get());
                    }
                }
                return make_array_vector(elements);
            }
            default:
                POLLUX_NYI(
                    "literalsToArrayVector not supported for type case '{}'", typeCase);
        }
    }

    core::TypedExprPtr SubstraitPolluxExprConverter::toPolluxExpr(
        const ::substrait::Expression::Cast &castExpr,
        const RowTypePtr &inputType) {
        auto type = substraitParser_.parseType(castExpr.type());
        bool nullOnFailure = isNullOnFailure(castExpr.failure_behavior());

        std::vector<core::TypedExprPtr> inputs{
            toPolluxExpr(castExpr.input(), inputType)
        };

        return std::make_shared<core::CastTypedExpr>(type, inputs, nullOnFailure);
    }

    core::TypedExprPtr SubstraitPolluxExprConverter::toPolluxExpr(
        const ::substrait::Expression &substraitExpr,
        const RowTypePtr &inputType) {
        core::TypedExprPtr polluxExpr;
        auto typeCase = substraitExpr.rex_type_case();
        switch (typeCase) {
            case ::substrait::Expression::RexTypeCase::kLiteral:
                return toPolluxExpr(substraitExpr.literal());
            case ::substrait::Expression::RexTypeCase::kScalarFunction:
                return toPolluxExpr(substraitExpr.scalar_function(), inputType);
            case ::substrait::Expression::RexTypeCase::kSelection:
                return toPolluxExpr(substraitExpr.selection(), inputType);
            case ::substrait::Expression::RexTypeCase::kCast:
                return toPolluxExpr(substraitExpr.cast(), inputType);
            case ::substrait::Expression::RexTypeCase::kIfThen:
                return toPolluxExpr(substraitExpr.if_then(), inputType);
            default:
                POLLUX_NYI(
                    "Substrait conversion not supported for Expression '{}'", typeCase);
        }
    }

    core::TypedExprPtr SubstraitPolluxExprConverter::toPolluxExpr(
        const ::substrait::Expression_IfThen &substraitIfThen,
        const RowTypePtr &inputType) {
        std::vector<core::TypedExprPtr> inputs;
        if (substraitIfThen.has_else_()) {
            inputs.reserve(substraitIfThen.ifs_size() * 2 + 1);
        } else {
            inputs.reserve(substraitIfThen.ifs_size() * 2);
        }

        TypePtr resultType;
        for (auto &ifExpr: substraitIfThen.ifs()) {
            auto ifClauseExpr = toPolluxExpr(ifExpr.if_(), inputType);
            inputs.emplace_back(ifClauseExpr);
            auto thenClauseExpr = toPolluxExpr(ifExpr.then(), inputType);
            inputs.emplace_back(thenClauseExpr);

            if (!thenClauseExpr->type()->containsUnknown()) {
                resultType = thenClauseExpr->type();
            }
        }

        if (substraitIfThen.has_else_()) {
            auto elseClauseExpr = toPolluxExpr(substraitIfThen.else_(), inputType);
            inputs.emplace_back(elseClauseExpr);
            if (!resultType && !elseClauseExpr->type()->containsUnknown()) {
                resultType = elseClauseExpr->type();
            }
        }

        POLLUX_CHECK_NOT_NULL(resultType, "Result type not found");

        return std::make_shared<const core::CallTypedExpr>(
            resultType, std::move(inputs), "if");
    }
} // namespace kumo::pollux::substrait
